from core.function import Function
from utils.functions_conv import im2col_array, col2im


class Im2col(Function):
    def __init__(self, kernel_size, stride, pad, to_matrix):
        """
        Im2col 操作的函数类，用于将输入的图像转换为用于卷积操作的矩阵形式。

        Args:
            kernel_size (int or tuple of int): 卷积核大小。
            stride (int or tuple of int): 步幅。
            pad (int or tuple of int): 填充。
            to_matrix (bool): 是否将结果转换为矩阵形式。
        """
        super().__init__()
        self.input_shape = None
        self.kernel_size = kernel_size
        self.stride = stride
        self.pad = pad
        self.to_matrix = to_matrix

    def forward(self, x):
        """
        前向传播，执行 Im2col 操作。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: 输出变量。
        """
        self.input_shape = x.shape
        y = im2col_array(x, self.kernel_size, self.stride, self.pad,
                         self.to_matrix)
        return y

    def backward(self, gy):
        """
        反向传播，计算输入变量的梯度。

        Args:
            gy (Variable): 输出变量的梯度。

        Returns:
            Variable: 输入变量的梯度。
        """
        gx = col2im(gy, self.input_shape, self.kernel_size, self.stride,
                    self.pad, self.to_matrix)
        return gx

    def im2col(self, x):
        """
        执行 Im2col 操作。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: 输出变量。
        """
        y = self(x)
        return y
